#encoding: utf-8
''' This module contains the handlers for parsing xml in extracting
    docs.
'''
import xml.sax.handler
from objects.term import Query
from general.langconv import Converter
from general.stringoperations import xml_escape, unify_punctuation

class QueryLoadHandler(xml.sax.ContentHandler):
    def __init__(self):
        self.buffer = ''
        self.query_dict = {}
        self.current_query = {}

    def load_query_types(self, file_path):
        self.query_types = {}
        for line in open(file_path):
            info = line.split('\t')
            self.query_types[info[0]] = info[1]

    def compare(self, q1, q2):
        if q1['name'] != q2['name'] or q1['docid'] != q2['docid']:
            print('Query Conflict:\n\t{0}\n\t{1}'.format(q1, q2))

    def startElement(self, name, attributes):
        if name == 'query':
            self.buffer = ''
            self.current_query = {}
            self.current_query['id'] = attributes['id']
    def characters(self, data):
        self.buffer += data
    def endElement(self, name):
        if name == 'name':
            self.current_query['name'] = self.buffer.strip()
            self.buffer = ''
        elif name == 'docid':
            self.current_query['docid'] = self.buffer.strip()
            self.buffer = ''
        elif name == 'query':
            if self.current_query['docid'].find('CMN') != -1 or \
                    self.current_query['docid'].find('cmn') != -1:
                self.current_query['entity_id'] = self.query_types[self.current_query['id']]
                if self.current_query['id'] not in self.query_dict:
                    self.query_dict[self.current_query['id']] = self.current_query
                else:
                    self.compare(self.current_query, self.query_dict[self.current_query['id']])

class QueryFileHandler(xml.sax.handler.ContentHandler):
    '''This handler is used to parse the query XML and then extract
       all doc id of Chinese query which is stored in setting.
    '''
    def __init__(self):
        #
        self.inDocid = False
        self.setting = set()

    def startElement(self, name, attributes):
        if name == "docid":
            self.buffer = ""
            self.inDocid = True

    def characters(self, data):
        if self.inDocid:
            self.buffer += data

    def endElement(self, name):
        if name == "docid":
            self.inDocid = False
            if self.buffer and (self.buffer.find('CMN') != -1 or
            self.buffer.find('cmn') != -1):
                self.setting.add(self.buffer)


class QueryDocHandler(xml.sax.handler.ContentHandler):
    '''This handler is used to deal with the query-related file to extract
       all the query-ralted doc, and to save each doc content in docDict
       indexed by doc id.
    '''
    def __init__(self, idList=[]):
        self.idList = idList    #contain docids we need
        self.inRecord = False
        self.docid = ""     #id of the current doc dealing with
        self.docDict = {}   #store the content of docids needed

    def startElement(self, name, attributes):
        if name == "DOC" and attributes['id'] in self.idList:
            self.docid = attributes['id']
            self.buffer = "<DOC id=\"%s\" type=\"%s\">" \
                    % (xml_escape(attributes['id']), xml_escape(attributes['type']))
            self.inRecord = True
        elif self.inRecord:
            self.buffer += "<%s>" % (name)
        else:
            pass

    def characters(self, data):
        if self.inRecord:
            self.buffer += xml_escape(data)

    def endElement(self, name):
        if self.inRecord:
            self.buffer += "</%s>" % (name)
        if name == "DOC" and self.docid != "":
            self.inRecord = False
            self.docDict[self.docid] = self.buffer  #set encode
            self.docid = ""


class QueryExtractionHandler(xml.sax.handler.ContentHandler):
    '''This handler is used to parse the query XML and then extract
       all Chinese query which is stored in query dict.
    '''
    def __init__(self, isDebug=False):
        self.query = ""
        self.inQuery = "Out"
        self.state = ""
        self.converter = Converter('zh-hans')
        self.dict = {}
        self.isDebug = isDebug

    def startElement(self, name, attributes):
        self.state = name
        if name == "query":
            self.query = Query(queryid=attributes['id'], isDebug=self.isDebug)
            self.inQuery = "Mid"
        if name != "query" and self.inQuery == "Mid":
            self.inQuery = "In"
            self.buffer = ""

    def characters(self, data):
        if self.inQuery == "In":
            self.buffer += data.strip()


    def endElement(self, name):
        if name == "query":
            self.inQuery = "Out"
            if self.query.docid.find('CMN') != -1 or\
               self.query.docid.find('cmn') != -1:
                if self.query.docid in self.dict:
                    for q in self.dict[self.query.docid]:
                        if q.queryid == self.query.queryid:
                            return
                    self.dict[self.query.docid].append(self.query)
                else:
                    self.dict[self.query.docid] = [self.query]
            self.query = ""
        elif self.inQuery == "In":
            if name == "beg" or name == "end":
                setattr(self.query, name, int(self.buffer))
            elif name == "name":
                simplified_name = self.converter.convert(self.buffer)
                self.query.name = unify_punctuation(simplified_name, isEscape=False)
                self.query.name_nospace = self.query.name.replace(' ', '')
                if self.query.name.find(' ') != -1:
                    self.query.name_with_newline = self.query.name.replace(' ', '\n')
                else:
                    self.query.name_with_newline = ""
            else:
                setattr(self.query, name, self.buffer)
            self.buffer = ""
